package com.wang.tx.mysql;

import com.alibaba.fastjson.JSONObject;
import com.wang.tx.core.TccCallback;
import com.wang.tx.core.TccTransaction;
import com.wang.tx.core.TccUtils;
import com.wang.utils.DatabaseUtils;

import javax.sql.DataSource;
import java.sql.*;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Executor;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;

import static com.wang.tx.core.TccCallback.*;

public class MysqlTccManager {
    private static Map<Class, Function<Object, List<TccTransaction>>> builder = new HashMap<>();
    private static final String create = "insert tcc_transaction_info(type, extra_info) value(?,?)";
    private static PreparedStatement createSql = null;
    public static <T> void register(Class<T> clazz, Function<T, List<TccTransaction>> func) {
        builder.put(clazz, (Function<Object, List<TccTransaction>>) func);
    }

    public static void init(DataSource dataSource) throws SQLException {
        Connection conn = dataSource.getConnection();
        createSql = conn.prepareStatement(create, Statement.RETURN_GENERATED_KEYS);
        MySqlTccCallback.updateSql = conn.prepareStatement(MySqlTccCallback.update);
        MySqlTccCallback.updateFailSql = conn.prepareStatement(MySqlTccCallback.updateFail);
    }

    public static void reRun(DataSource dataSource, Executor executor) throws SQLException, ClassNotFoundException {
        reRun(dataSource, executor, Throwable::printStackTrace);
    }

    // 重启时检查未完成事务
    public static void reRun(DataSource dataSource, Executor executor, Consumer<Exception> exceptionCallback) throws SQLException, ClassNotFoundException {
        List<TccTransactionInfo> list = DatabaseUtils.INSTANCE.select(dataSource.getConnection(),
                "select * from tcc_transaction_info where state in (0,1,2)", TccTransactionInfo.class);
        for (TccTransactionInfo info : list) {
            Class clazz = Class.forName(info.getType());
            Function<Object, List<TccTransaction>> func = builder.get(clazz);

            executor.execute(() -> {
                Object extraInfo = JSONObject.parseObject(info.getExtraInfo(), clazz);
                List<TccTransaction> txList = func.apply(extraInfo);
                List<TccTransaction> tryList = null;
                if (info.getState() == TRY_FAIL) {
                    List<Long> idList = JSONObject.parseArray(info.getHasTryList()).toJavaList(Long.class);
                    tryList = txList.stream().filter(i -> idList.contains(i.id())).collect(Collectors.toList());
                } else if (info.getState() == TRY_OK) {
                    tryList = txList;
                }
                try {
                    TccUtils.run(info.getState(), txList, tryList, new MySqlTccCallback(info.getId()));
                } catch (Exception e) {
                    exceptionCallback.accept(e);
                }
            });
        }
    }

    public static <T> boolean execute(T info) throws Exception {
        Function<T, List<TccTransaction>> func = (Function<T, List<TccTransaction>>) builder.get(info.getClass());
        List<TccTransaction> txList = func.apply(info);
        System.out.println(JSONObject.toJSONString(info));
        // 创建数据库记录
        createSql.setString(1, info.getClass().getName());
        createSql.setString(2, JSONObject.toJSONString(info));
        createSql.executeUpdate();
        ResultSet rs = createSql.getGeneratedKeys();
        rs.next();
        long id = rs.getLong(1);
        return TccUtils.run(INIT, txList, null, new MySqlTccCallback(id));
    }
}
